# DuckDB客户端模块 - 用于实时数据的持久化存储

import duckdb
import pandas as pd
import json
import logging
from typing import Dict, List, Any, Optional, Union
from datetime import datetime, timedelta
import os
from src.utils.schemas import STOCK_PRICES_SCHEMA, TECHNICAL_INDICATORS_SCHEMA, STOCK_LIST_SCHEMA

logger = logging.getLogger(__name__)

class DuckDBClient:
    """
    DuckDB客户端类，用于实时数据的持久化存储
    将Redis中的短期数据持久化到DuckDB中
    """
    
    def __init__(self, db_path: str = 'data/realtime.duckdb'):
        """
        初始化DuckDB连接
        
        Args:
            db_path: 数据库文件路径
        """
        self._connected = False
        try:
            # 确保目录存在
            os.makedirs(os.path.dirname(db_path), exist_ok=True)
            
            # 连接到DuckDB
            self.conn = duckdb.connect(db_path)
            self.test_connection()
            self._connected = True
            
            # 初始化表结构
            self._init_tables()
            
            logger.info(f"DuckDB连接初始化成功: {db_path}")
        except Exception as e:
            logger.error(f"DuckDB连接初始化失败: {str(e)}")
            raise

    def test_connection(self) -> bool:
        """测试DuckDB连接"""
        try:
            self.conn.execute('SELECT 1')
            self._connected = True
            return True
        except:
            self._connected = False
            return False
    
    def is_connected(self) -> bool:
        """获取连接状态"""
        if not self._connected:
            return self.test_connection()
        return self._connected
    
    def _init_tables(self):
        """
        初始化必要的表结构
        """
        try:
            # 创建股票行情表
            stock_prices_columns = [
                f"{col} {dtype.replace('SERIAL', 'INTEGER')}" for col, dtype in STOCK_PRICES_SCHEMA['columns'].items()
            ]
            self.conn.execute(f"""
                CREATE TABLE IF NOT EXISTS stock_prices (
                    {', '.join(stock_prices_columns)}
                );
            """)

            # 创建技术指标扩展表
            technical_indicators_columns = [
                f"{col} {dtype.replace('SERIAL', 'INTEGER')}" for col, dtype in TECHNICAL_INDICATORS_SCHEMA['columns'].items()
            ]
            technical_indicators_primary_key = f"PRIMARY KEY ({', '.join(TECHNICAL_INDICATORS_SCHEMA['primary_key'])})"
            technical_indicators_foreign_keys = "".join(
                [f", FOREIGN KEY ({fk['column']}) REFERENCES {fk['references']}" for fk in TECHNICAL_INDICATORS_SCHEMA.get('foreign_keys', [])]
            )
            self.conn.execute(f"""
                CREATE TABLE IF NOT EXISTS technical_indicators (
                    {', '.join(technical_indicators_columns)},
                    {technical_indicators_primary_key},
                    {technical_indicators_foreign_keys}
                );
            """)

            # 创建股票列表表
            stock_list_columns = [
                f"{col} {dtype.replace('SERIAL', 'INTEGER')}" for col, dtype in STOCK_LIST_SCHEMA['columns'].items()
            ]
            self.conn.execute(f"""
                CREATE TABLE IF NOT EXISTS stock_list (
                    {', '.join(stock_list_columns)}
                );
            """)

            logger.info("DuckDB表结构初始化成功")
        except Exception as e:
            logger.error(f"DuckDB表结构初始化失败: {str(e)}")
            raise
    
    def persist_realtime_quotes(self, quotes_data: List[Dict]):
        """
        持久化实时行情数据
        
        Args:
            quotes_data: 实时行情数据列表
        """
        try:
            if not quotes_data:
                return
            
            # 转换为DataFrame
            df = pd.DataFrame(quotes_data)
            
            # 确保时间戳列是datetime类型
            df['timestamp'] = pd.to_datetime(df['timestamp'])
            
            # 提取必要的列，其余放入JSON
            required_columns = ['symbol', 'timestamp', 'price', 'open', 'high', 'low', 'volume', 'amount', 
                               'bid_price', 'ask_price', 'bid_volume', 'ask_volume']
            
            # 创建data列，包含所有原始数据
            df['data'] = df.apply(lambda row: json.dumps({col: row[col] for col in row.index if col not in required_columns}), axis=1)
            
            # 只保留必要的列
            for col in required_columns:
                if col not in df.columns:
                    df[col] = None
            
            # 插入数据
            self.conn.execute("""
                INSERT OR REPLACE INTO realtime_quotes 
                SELECT * FROM df
            """)
            
            logger.info(f"成功持久化{len(df)}条实时行情数据")
        except Exception as e:
            logger.error(f"持久化实时行情数据失败: {str(e)}")

    def query_realtime_quotes(self, symbol: str = None, start_time: str = None, end_time: str = None) -> pd.DataFrame:
        """
        查询实时行情数据
        
        Args:
            symbol: 股票代码，如果为None则查询所有股票
            start_time: 开始时间，格式'YYYYMMDD HH:MM:SS'
            end_time: 结束时间，格式'YYYYMMDD HH:MM:SS'
            
        Returns:
            pd.DataFrame: 包含实时行情数据的DataFrame
        """
        try:
            query = "SELECT * FROM realtime_quotes"
            conditions = []
            params = []
            
            if symbol:
                conditions.append("symbol = ?")
                params.append(symbol)
            
            if start_time:
                conditions.append("timestamp >= ?")
                params.append(start_time)
            
            if end_time:
                conditions.append("timestamp <= ?")
                params.append(end_time)
            
            if conditions:
                query += " WHERE " + " AND ".join(conditions)
            
            query += " ORDER BY timestamp DESC"
            
            # 执行查询
            result = self.conn.execute(query, params).fetchdf()
            return result
        except Exception as e:
            logger.error(f"查询实时行情数据失败: {str(e)}")
            return pd.DataFrame()

    def close(self):
        """
        关闭DuckDB连接
        """
        if hasattr(self, 'conn'):
            self.conn.close()
            logger.info("DuckDB连接已关闭")
    
    def __del__(self):
        """
        析构函数，确保连接被关闭
        """
        self.close()